import time
import cv2
import numpy as np
import pytesseract
import win32gui
import win32con
import win32ui
import ctypes
from tkinter import ttk, messagebox, scrolledtext
import tkinter as tk
import threading
import schedule
import os
from datetime import datetime

# 设置 tesseract 路径（根据你的安装路径修改）
pytesseract.pytesseract.tesseract_cmd = r'H:\Tesseract-OCR\tesseract.exe'

# ==================== 配置 ====================
SAVE_DIR = "results"
os.makedirs(SAVE_DIR, exist_ok=True)
RESULT_FILE = os.path.join(SAVE_DIR, "numbers.csv")

# 全局变量
window_title = ""
target_num1 = ""
target_num2 = ""
search_region = None  # 用户输入的搜索区域 (x, y, w, h) 相对于屏幕
found_rect1 = None    # OCR 找到的数字1区域 (x, y, w, h) 相对于窗口
found_rect2 = None    # OCR 找到的数字2区域
is_monitoring = False

# ==================== 工具函数 ====================

def find_window_fuzzy(partial_title):
    """模糊查找窗口（不区分大小写）"""
    def callback(hwnd, results):
        if win32gui.IsWindowVisible(hwnd):
            title = win32gui.GetWindowText(hwnd)
            if partial_title.lower() in title.lower():
                results.append(hwnd)
        return True
    results = []
    win32gui.EnumWindows(callback, results)
    return results[0] if results else None

def capture_window_advanced(hwnd):
    """使用 PrintWindow 截图（支持最小化）"""
    left, top, right, bottom = win32gui.GetWindowRect(hwnd)
    width = right - left
    height = bottom - top

    hwnd_dc = win32gui.GetWindowDC(hwnd)
    mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
    save_dc = mfc_dc.CreateCompatibleDC()

    save_bitmap = win32ui.CreateBitmap()
    save_bitmap.CreateCompatibleBitmap(mfc_dc, width, height)
    save_dc.SelectObject(save_bitmap)

    result = ctypes.windll.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 0x00000002)
    if result == 0:
        save_dc.BitBlt((0, 0), (width, height), mfc_dc, (0, 0), win32con.SRCCOPY)

    signed_ints = save_bitmap.GetBitmapBits(True)
    img = np.frombuffer(signed_ints, dtype='uint8')
    img.shape = (height, width, 4)
    img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

    win32gui.DeleteObject(save_bitmap.GetHandle())
    save_dc.DeleteDC()
    mfc_dc.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwnd_dc)

    return img, (left, top)

def find_numbers_in_region(img, search_x, search_y, search_w, search_h, num1, num2):
    """
    在指定区域内 OCR，查找两个数字的位置
    返回：(rect1, rect2) 坐标相对于窗口
    """
    # 提取搜索区域
    h, w = img.shape[:2]
    roi_x = max(0, search_x)
    roi_y = max(0, search_y)
    roi_w = min(search_w, w - roi_x)
    roi_h = min(search_h, h - roi_y)

    roi = img[roi_y:roi_y+roi_h, roi_x:roi_x+roi_w]

    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    data = pytesseract.image_to_data(
        thresh, config='--psm 6', output_type=pytesseract.Output.DICT)
    n_boxes = len(data['text'])

    rect1 = rect2 = None

    for i in range(n_boxes):
        text = data['text'][i].strip()
        if not text.isdigit():
            continue

        x = data['left'][i] + roi_x
        y = data['top'][i] + roi_y
        w = data['width'][i]
        h = data['height'][i]

        if num1 in text and rect1 is None:
            rect1 = (x, y, w, h)
        if num2 in text and rect2 is None:
            rect2 = (x, y, w, h)

    return rect1, rect2

def ocr_at_rect(img, rect):
    """在指定矩形内 OCR 识别数字"""
    x, y, w, h = rect
    roi = img[y:y+h, x:x+w]
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    config = '--psm 8 -c tessedit_char_whitelist=0123456789'
    text = pytesseract.image_to_string(thresh, config=config).strip()
    return text


def save_result(text1, text2):
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(RESULT_FILE, 'a', encoding='utf-8') as f:
        if os.path.getsize(RESULT_FILE) == 0:
            f.write("Time,Value1,Value2\n")
        f.write(f"{now},{text1},{text2}\n")
    print(f"[{now}] 识别结果: {text1} | {text2}")


def monitoring_job():
    """定时任务：在已定位的区域内识别"""
    global window_title, found_rect1, found_rect2, search_region

    hwnd = find_window_fuzzy(window_title)
    if not hwnd:
        print("❌ 窗口丢失，跳过本次识别")
        return

    img, _ = capture_window_advanced(hwnd)
    if img is None:
        return

    text1 = text2 = "N/A"
    if found_rect1:
        text1 = ocr_at_rect(img, found_rect1)
    if found_rect2:
        text2 = ocr_at_rect(img, found_rect2)

    save_result(text1, text2)

def start_monitoring():
    """监控线程"""
    global is_monitoring
    schedule.every(5).minutes.do(monitoring_job)
    while is_monitoring:
        schedule.run_pending()
        time.sleep(1)


# ==================== GUI 主界面 ====================

class OCRMonitorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("OCR数字定位监控器 v2.0")
        self.root.geometry("600x500")
        self.root.resizable(False, False)

        self.create_widgets()

    def create_widgets(self):
        frame = ttk.Frame(self.root, padding="10")
        frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

        row = 0
        # 窗口标题
        ttk.Label(frame, text="窗口标题:").grid(row=row, column=0, sticky=tk.W, pady=3)
        self.entry_title = ttk.Entry(frame, width=40)
        self.entry_title.grid(row=row, column=1, pady=3)

        row += 1
        # 数字1
        ttk.Label(frame, text="数字目标 1:").grid(row=row, column=0, sticky=tk.W, pady=3)
        self.entry_num1 = ttk.Entry(frame, width=40)
        self.entry_num1.grid(row=row, column=1, pady=3)

        row += 1
        # 数字2
        ttk.Label(frame, text="数字目标 2:").grid(row=row, column=0, sticky=tk.W, pady=3)
        self.entry_num2 = ttk.Entry(frame, width=40)
        self.entry_num2.grid(row=row, column=1, pady=3)

        row += 1
        # 搜索区域
        ttk.Label(frame, text="搜索区域 X:").grid(row=row, column=0, sticky=tk.W, pady=3)
        self.entry_x = ttk.Entry(frame, width=10)
        self.entry_x.grid(row=row, column=1, sticky=tk.W, padx=(0,5))
        ttk.Label(frame, text="Y:").grid(row=row, column=1, sticky=tk.W, padx=(80,0))
        self.entry_y = ttk.Entry(frame, width=10)
        self.entry_y.grid(row=row, column=1, sticky=tk.W, padx=(100,0))

        row += 1
        ttk.Label(frame, text="Width:").grid(row=row, column=0, sticky=tk.W, pady=3)
        self.entry_w = ttk.Entry(frame, width=10)
        self.entry_w.grid(row=row, column=1, sticky=tk.W, padx=(0,5))
        ttk.Label(frame, text="Height:").grid(row=row, column=1, sticky=tk.W, padx=(80,0))
        self.entry_h = ttk.Entry(frame, width=10)
        self.entry_h.grid(row=row, column=1, sticky=tk.W, padx=(120,0))

        row += 1
        # 按钮
        btn_frame = ttk.Frame(frame)
        btn_frame.grid(row=row, column=0, columnspan=2, pady=10)
        ttk.Button(btn_frame, text="开始定位并监控", command=self.start).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="停止监控", command=self.stop).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="立即识别一次", command=self.quick_ocr).pack(side=tk.LEFT, padx=5)

        row += 1
        # 日志
        ttk.Label(frame, text="状态日志:").grid(row=row, column=0, sticky=tk.W, pady=5)
        self.log_text = scrolledtext.ScrolledText(frame, height=12, width=70)
        self.log_text.grid(row=row+1, column=0, columnspan=2, pady=5)

    def log(self, msg):
        self.log_text.insert(tk.END, msg + "\n")
        self.log_text.see(tk.END)

    def start(self):
        global window_title, target_num1, target_num2, search_region, found_rect1, found_rect2, is_monitoring

        window_title = self.entry_title.get().strip()
        target_num1 = self.entry_num1.get().strip()
        target_num2 = self.entry_num2.get().strip()

        try:
            x = int(self.entry_x.get())
            y = int(self.entry_y.get())
            w = int(self.entry_w.get())
            h = int(self.entry_h.get())
            search_region = (x, y, w, h)
        except ValueError:
            messagebox.showerror("错误", "X, Y, Width, Height 必须是整数！")
            return

        if not window_title:
            messagebox.showwarning("警告", "请输入窗口标题！")
            return
        if not target_num1 and not target_num2:
            messagebox.showwarning("警告", "至少输入一个识别数字！")
            return

        self.log(f"🔍 查找窗口: '{window_title}'")

        hwnd = find_window_fuzzy(window_title)
        if not hwnd:
            messagebox.showerror("错误", "未找到匹配的窗口！")
            return

        img, (win_x, win_y) = capture_window_advanced(hwnd)
        if img is None:
            messagebox.showerror("错误", "截图失败！")
            return

        # 转换搜索区域为相对于窗口的坐标
        rel_x = x - win_x
        rel_y = y - win_y

        self.log(f"🔍 在区域 ({rel_x}, {rel_y}, {w}, {h}) 内搜索数字...")

        rect1, rect2 = find_numbers_in_region(
            img, rel_x, rel_y, w, h, target_num1, target_num2
        )

        if rect1:
            self.log(f"✅ 找到 '{target_num1}' 区域: {rect1}")
            found_rect1 = rect1
        else:
            self.log(f"❌ 未找到数字 '{target_num1}'")

        if rect2:
            self.log(f"✅ 找到 '{target_num2}' 区域: {rect2}")
            found_rect2 = rect2
        else:
            self.log(f"❌ 未找到数字 '{target_num2}'")

        if not found_rect1 and not found_rect2:
            messagebox.showerror("错误", "未能在指定区域内找到任一目标数字！")
            return

        # 开始监控
        is_monitoring = True
        thread = threading.Thread(target=start_monitoring, daemon=True)
        thread.start()
        self.log("🟢 监控已启动，每5分钟识别一次...")

    def stop(self):
        global is_monitoring
        is_monitoring = False
        self.log("🔴 监控已停止")

    def quick_ocr(self):
        global window_title, found_rect1, found_rect2
        if not window_title:
            messagebox.showwarning("警告", "请先启动监控！")
            return

        hwnd = find_window_fuzzy(window_title)
        if not hwnd:
            self.log("❌ 窗口未找到")
            return

        img, _ = capture_window_advanced(hwnd)
        if img is None:
            self.log("📷 截图失败")
            return

        text1 = ocr_at_rect(img, found_rect1) if found_rect1 else "N/A"
        text2 = ocr_at_rect(img, found_rect2) if found_rect2 else "N/A"
        self.log(f"📌 快速识别: {text1} | {text2}")


# ==================== 主程序 ====================
def main():
    root = tk.Tk()
    app = OCRMonitorApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()